#!/usr/bin/env python  
#-*- coding:utf-8 _*-  
""" 
@author:hello_life 
@license: Apache Licence 
@file: clue_dataset.py 
@time: 2022/04/30
@software: PyCharm 
description:
"""
import numpy as np
import torch
from torch.utils.data import Dataset

from utils.data_utils import content2id,label2id,mask_make

class CLUE_Dataset(Dataset):
    def __init__(self,config,data_dir):
        self.config=config
        self.data=np.load(data_dir,allow_pickle=True)

    def __len__(self):
        assert len(self.data["content"])==len(self.data["label"])
        return len(self.data["content"])

    def __getitem__(self, idx):
        x=self.data["content"][idx]
        y=self.data["label"][idx]

        if self.config.do_mask:
            mask=mask_make(x,self.config)
            mask=self.mask_transformer(mask)

        x=self.x_transformer(x)
        y=self.y_transformer(y)

        if self.config.do_mask:
            return x,mask,y
        else:
            return x,y

    def x_transformer(self,x):
        x=content2id(x,self.config)
        return torch.tensor(x,dtype=torch.int32).to(self.config.device)

    def y_transformer(self,y):
        y=label2id(y,self.config)
        return torch.tensor(y,dtype=torch.long).to(self.config.device)

    def mask_transformer(self,mask):
        return torch.tensor(mask,dtype=torch.bool).to(self.config.device)
